#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户表字段扩展迁移脚本
添加新的用户管理字段
"""

import sqlite3
import os
from datetime import datetime

def migrate_database():
    """执行数据库迁移"""
    # 数据库文件路径 - 使用instance目录下的数据库文件
    db_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'instance', 'dispatch_system.db')
    print(f"使用数据库文件: {db_path}")
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        print("开始执行用户表字段扩展迁移...")
        
        # 添加新字段的SQL语句
        alter_statements = [
            "ALTER TABLE users ADD COLUMN area TEXT",
            "ALTER TABLE users ADD COLUMN id_card TEXT", 
            "ALTER TABLE users ADD COLUMN ethnicity TEXT",
            "ALTER TABLE users ADD COLUMN gender TEXT",
            "ALTER TABLE users ADD COLUMN hire_date DATE",
            "ALTER TABLE users ADD COLUMN is_driver BOOLEAN DEFAULT 0",
            "ALTER TABLE users ADD COLUMN password_field TEXT",
            "ALTER TABLE users ADD COLUMN status TEXT DEFAULT 'active'",
            "ALTER TABLE users ADD COLUMN resignation_date DATE",
            "ALTER TABLE users ADD COLUMN electrician_cert_expiry DATE",
            "ALTER TABLE users ADD COLUMN height_cert_expiry DATE",
            "ALTER TABLE users ADD COLUMN batch_number TEXT",
            "ALTER TABLE users ADD COLUMN remarks TEXT"
        ]
        
        # 执行每个ALTER语句
        for statement in alter_statements:
            try:
                cursor.execute(statement)
                field_name = statement.split('ADD COLUMN ')[1].split(' ')[0]
                print(f"✓ 成功添加字段: {field_name}")
            except sqlite3.OperationalError as e:
                if "duplicate column name" in str(e).lower():
                    field_name = statement.split('ADD COLUMN ')[1].split(' ')[0]
                    print(f"- 字段已存在，跳过: {field_name}")
                else:
                    print(f"✗ 添加字段失败: {e}")
                    raise
        
        # 更新现有用户的必填字段默认值
        update_statements = [
            "UPDATE users SET area = '未分配' WHERE area IS NULL",
            "UPDATE users SET id_card = '' WHERE id_card IS NULL",
            "UPDATE users SET ethnicity = '汉族' WHERE ethnicity IS NULL",
            "UPDATE users SET gender = 'male' WHERE gender IS NULL",
            "UPDATE users SET hire_date = date('now') WHERE hire_date IS NULL",
            "UPDATE users SET is_driver = 0 WHERE is_driver IS NULL",
            "UPDATE users SET status = 'active' WHERE status IS NULL"
        ]
        
        for statement in update_statements:
            cursor.execute(statement)
            print(f"✓ 更新默认值: {statement.split('SET ')[1].split(' =')[0]}")
        
        # 提交更改
        conn.commit()
        print("\n✓ 用户表字段扩展迁移完成！")
        
        # 验证新字段
        cursor.execute("PRAGMA table_info(users)")
        columns = cursor.fetchall()
        print("\n当前用户表字段:")
        for col in columns:
            print(f"  - {col[1]} ({col[2]})")
            
    except Exception as e:
        print(f"✗ 迁移失败: {e}")
        conn.rollback()
        raise
    finally:
        conn.close()

if __name__ == "__main__":
    migrate_database()